package com.smh.aiagent.chatmemory;

import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.smh.aiagent.mapper.UserMessageMapper;
import com.smh.aiagent.model.pojo.UserMessage;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;

/**
 * @author shiminghui
 * @date 2025/5/31 16:49
 * @description: TODO
 */
@Component
public class DataBasedChatMemory implements ChatMemory {


    @Autowired
    private UserMessageMapper userMessageMapper;


    @Override
    public void add(String conversationId, Message message) {
        saveToDataBase(conversationId, List.of(message));
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        saveToDataBase(conversationId, messages);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        return getMessages(conversationId, lastN);
    }

    @Override
    public void clear(String conversationId) {
        userMessageMapper.delete(new LambdaQueryWrapper<UserMessage>()
                .eq(UserMessage::getConversationId, conversationId));
    }


    public List<Message> getMessages(String conversationId, int lastN) {
        List<UserMessage> userMessages = userMessageMapper.selectMessagesByConversationIdAndLastN(conversationId, lastN);
        ArrayList<Message> messages = new ArrayList<>();
        for (UserMessage userMessage : userMessages) {
            Message message;
            switch (userMessage.getMessageType()) {
                case USER -> message = new org.springframework.ai.chat.messages.UserMessage(userMessage.getContent());
                case ASSISTANT -> message = new AssistantMessage(userMessage.getContent());
                case SYSTEM -> message = new SystemMessage(userMessage.getContent());
                case TOOL -> {
                    String content = userMessage.getContent();
                    JSONArray jsonArray = JSONUtil.parseArray(content);
                    List<ToolResponseMessage.ToolResponse> responses = new ArrayList<>();
                    for (Object obj : jsonArray) {
                        ToolResponseMessage.ToolResponse toolResponse = JSONUtil.toBean((JSONObject) obj, ToolResponseMessage.ToolResponse.class);
                        responses.add(toolResponse);
                    }
                    message = new ToolResponseMessage(responses);
                }
                default -> throw new IllegalArgumentException("Invalid message type: " + userMessage.getMessageType());
            }
            messages.addFirst(message); // 将消息添加到列表的开头,使倒序变为正序
        }
        return messages;
    }

    public void saveToDataBase(String conversationId, List<Message> messages) {
        // 将消息转化为对应的子消息对象
        ArrayList<UserMessage> userMessages = new ArrayList<>();

        for (Message message : messages) {
            UserMessage userMessage = new UserMessage();
            switch (message.getMessageType()) {
                case USER -> userMessage.setMessageType(UserMessage.UserMessageType.USER);
                case ASSISTANT -> userMessage.setMessageType(UserMessage.UserMessageType.ASSISTANT);
                case SYSTEM -> userMessage.setMessageType(UserMessage.UserMessageType.SYSTEM);
                case TOOL -> userMessage.setMessageType(UserMessage.UserMessageType.TOOL);
                default -> throw new IllegalArgumentException("Invalid message type: " + message.getMessageType());
            }
            userMessage.setConversationId(conversationId);
            if (message.getMessageType().equals(MessageType.TOOL)) {
                ToolResponseMessage toolResponseMessage = (ToolResponseMessage) message;
                List<ToolResponseMessage.ToolResponse> responses = toolResponseMessage.getResponses();
                String jsonStr = JSONUtil.toJsonStr(responses);
                userMessage.setContent(jsonStr);
            } else {
                userMessage.setContent(message.getText());
            }
            userMessage.setCreateTime(LocalDateTime.now());
            userMessages.add(userMessage);
        }
        //  保存到数据库
        userMessageMapper.insert(userMessages);
    }


}
